package com.edata.bigdata.postgres
import com.edata.bigdata.annotations.Edata_Saver
import com.edata.bigdata.basic.Saver
import com.edata.bigdata.utils.{DataFrameUtils, QueryLangUtils}
import org.apache.spark.sql.types.StructField
import org.apache.spark.sql.{DataFrame, Row, SaveMode, SparkSession}

import java.sql.{Connection, DriverManager, Statement}
import java.util.Properties
import scala.collection.mutable.ArrayBuffer
@Edata_Saver(target = "PGSAVER")
class PgSaver extends Saver with PgConnector {
  override var session: SparkSession = _
  override var sourceType: String = "PG"

  def createConnection(): Connection = {
    val props = new Properties()
    props.put("user", JDBC_USER)
    props.put("password", JDBC_PASSWORD)
    props.setProperty("useSSL", "false")
    props.setProperty("useUnicode", "true")
    props.setProperty("characterEncoding", "utf-8")
    val URL = s"${JDBC_PREFIX}${JDBC_IP}:${JDBC_PORT}/${JDBC_DATABASE}?${JDBC_BATCHPARAMS}"
    DriverManager.getConnection(URL, props)
  }

  override def save(data: DataFrame, args: String*): Boolean = {
    try{
      data.write.format("jdbc").mode(SaveMode.Append).save()
      return true
    }catch {
      case e:Exception => {
        LOGGER.error(s"save data fail:${e.getMessage}")
        return false
      }
    }
  }

  override def update(data: DataFrame, args: String*): Boolean = {
    var connection: Connection = null
    data.foreachPartition((iter: Iterator[Row]) => {
      try {
        connection = createConnection()
        val st: Statement = connection.createStatement()
        iter.foreach((row: Row) => {
          val cv_array = new ArrayBuffer[String]()
          var id_condition=""
          row.schema.foreach((c: StructField) => {
            if (args(1) != c.name) {
              cv_array.append(s"${c.name}=${QueryLangUtils.convertToQueryStr(DataFrameUtils.getValueByCol(row, c.name, c.dataType.typeName),c.dataType.typeName)}")
            }else{
              id_condition=s"${c.name}=${QueryLangUtils.convertToQueryStr(DataFrameUtils.getValueByCol(row, c.name, c.dataType.typeName),c.dataType.typeName)}"
            }
          })
          val sql = s"update ${args(0)} set (${cv_array.toArray.mkString(",")}) where ${id_condition}"
          LOGGER.info(s"Executing SQL:${sql}")
          st.addBatch(sql)
        })
        //调用st.executeBatch()相当于告诉编译器每个分区都需要返回Array[Int]，但是在foreachPartition中只能返回Unit
        val result = st.executeBatch()
      } catch {
        case e: Exception => e.printStackTrace()
      } finally {
        if (connection != null) {
          connection.close()
        }
      }
    })
    return true
  }

  override def delete(data: DataFrame, args: String*): Boolean = {
    var conn: Connection = null
    data.foreachPartition((iter: Iterator[Row]) => {
      try {
        conn = createConnection()
        val st: Statement = conn.createStatement()
        iter.foreach((row: Row) => {
          val cv_array = new ArrayBuffer[String]()
          row.schema.foreach((c: StructField) => {
            cv_array.append(s"${c.name}=${QueryLangUtils.convertToQueryStr(DataFrameUtils.getValueByCol(row, c.name, c.dataType.typeName),c.dataType.typeName)}")
          })
          val sql = s"delete from ${args(0)} where (${cv_array.toArray.mkString(",")})"
          LOGGER.info(s"Executing SQL:${sql}")
          st.addBatch(sql)
        })
        //调用st.executeBatch()相当于告诉编译器每个分区都需要返回Array[Int]，但是在foreachPartition中只能返回Unit
        val result = st.executeBatch()
      } catch {
        case e: Exception => e.printStackTrace()
      } finally {
        if (conn != null) {
          conn.close()
        }
      }
    })
    return true
  }

  override def saveOrUpdate(data: DataFrame, args: String*): Boolean = {
    var connection: Connection = null
    data.foreachPartition((iter: Iterator[Row]) => {
      try {
        connection = createConnection()
        val st: Statement = connection.createStatement()
        iter.foreach((row: Row) => {
          val c_array = new ArrayBuffer[String]()
          val v_array = new ArrayBuffer[String]()
          val cv_array = new ArrayBuffer[String]()
          row.schema.foreach((c: StructField) => {
            c_array.append(c.name)
            v_array.append(s"${QueryLangUtils.convertToQueryStr(DataFrameUtils.getValueByCol(row,c.name,c.dataType.typeName),c.dataType.typeName)}")
            if (args(1)!=c.name) {
              cv_array.append(s"${c.name}=${QueryLangUtils.convertToQueryStr(DataFrameUtils.getValueByCol(row,c.name,c.dataType.typeName),c.dataType.typeName)}")
            }
          })
          val sql = s"INSERT INTO ${args(0)} (${c_array.toArray.mkString(",")}) VALUES (${v_array.toArray.mkString(",")}) ON CONFLICT(${args(1)}) DO UPDATE SET ${cv_array.toArray.mkString(",")}"
          LOGGER.info(s"Executing SQL:${sql}")
          st.addBatch(sql)
        })
        //调用st.executeBatch()相当于告诉编译器每个分区都需要返回Array[Int]，但是在foreachPartition中只能返回Unit
        val result = st.executeBatch()
      } catch {
        case e: Exception => e.printStackTrace()
      } finally {
        if (connection != null) {
          connection.close()
        }
      }
    })
    return true
  }
}
